import https from 'https';
import { prisma } from '../service/mysql.js';
import { logger } from '../tools/index.js';
const apiKey = process.env.DASHSCOPE_API_KEY;
const PATH_MAPPING = {
    TEXT_GENERATION_API: '/api/v1/services/aigc/text-generation/generation',
    MULTIMODAL_GENERATION_API: '/api/v1/services/aigc/multimodal-generation/generation'
}

export const chatSSE = async function (prompt, callback = () => { }, context) {
    return await chatBasicSSE({
        model: "qwen-long",
        input: {
            messages: [
                {
                    role: "system",
                    content: "You are a helpful assistant."
                },
                {
                    role: "user",
                    content: prompt
                }
            ]
        },
        parameters: {
            temperature: 0.1,
            top_p: 0.1,
            result_format: "message",
            incremental_output: true
        }
    }, PATH_MAPPING.TEXT_GENERATION_API, callback, context).then(res => {
        if (res?.output?.incrementalText) {
            res.output.incrementalText = undefined;
        }
        return res;
    });
}

export const chatVlSSE = async function (imageUrl, prompt, callback = () => { }, context) {
    return await chatBasicSSE({
        model: "qwen-vl-chat-v1",
        input: {
            messages: [
                {
                    role: "system",
                    content: [
                        { text: "You are a helpful assistant." }
                    ]
                },
                {
                    role: "user",
                    content: [
                        { image: imageUrl },
                        { text: prompt }
                    ]
                }
            ]
        },
        parameters: {
            temperature: 0.1,
            top_p: 0.1,
            result_format: "message",
            incremental_output: true
        }
    }, PATH_MAPPING.MULTIMODAL_GENERATION_API, callback, context).then(res => {
        if (res?.output?.text && res?.output?.incrementalText) {
            res.output.text = res.output.incrementalText;
            res.output.incrementalText = undefined;
        }
        return res;
    });
}

const chatBasicSSE = async (params, path, callback = () => { }, context) => {
    const data = JSON.stringify(params);
    const messageContent = params.input.messages[params.input.messages.length - 1].content;
    const prompt = Array.isArray(messageContent) ? JSON.stringify(messageContent) : messageContent;
    logger.debug("chat qwen start, prompt = %s", prompt);
    const options = {
        hostname: 'dashscope.aliyuncs.com',
        path: path,
        method: 'POST',
        headers: {
            'Authorization': `Bearer ${apiKey}`,
            'Content-Type': 'application/json',
            'X-DashScope-SSE': 'enable',
            'Content-Length': Buffer.byteLength(data),
            // qwen-long 不支持
            // 'X-DashScope-Plugin': '{"calculator":{}}'
        }
    };

    return new Promise((resolve, reject) => {
        let resData = {
            output: {
                text: '',
                incrementalText: '',
                finish_reason: null
            },
            usage: {
                total_tokens: 0,
                input_tokens: 0,
                output_tokens: 0,
            },
            err: '',
            request_id: null
        };
        const req = https.request(options, async (res) => {
            let buffer = '';
            res.on('data', async (chunk) => {
                buffer += chunk.toString();

                // 分割收到的数据段，因为可能包含多个完整的SSE消息
                let messages = buffer.split(/\n\n/);
                buffer = messages.pop(); // 保留未完成的部分

                messages.forEach(async msg => {
                    msg.split(/\n/).forEach(async itemMsg => {
                        if (itemMsg.startsWith('data:')) {
                            const dataPart = itemMsg.replace('data:', '').trim();
                            const dataPartJson = JSON.parse(dataPart);
                            if (dataPartJson.code && dataPartJson.message) {
                                resData.request_id = dataPartJson.request_id;
                                resData.err = dataPartJson.message;
                                await prisma.t_llm_log.create({
                                    data: {
                                        gmt_create: new Date(),
                                        input_tokens: resData.usage.input_tokens,
                                        output_tokens: resData.usage.output_tokens,
                                        total_tokens: resData.usage.total_tokens,
                                        prompt: prompt,
                                        content: resData.output.text,
                                        response: JSON.stringify(resData),
                                        status: 'failed',
                                        error: resData.err,
                                        request_id: resData.request_id,
                                        chat_request_id: context.requestId
                                    }
                                });
                                callback('error', dataPartJson.message);
                                reject(dataPartJson.message);
                            } else {
                                const incrementalText = dataPartJson.output.choices[0].message.content;
                                resData = {
                                    output: {
                                        incrementalText: incrementalText,
                                        text: resData.output.text + incrementalText,
                                        finish_reason: dataPartJson.output.choices[0].finish_reason
                                    },
                                    usage: dataPartJson.usage,
                                    request_id: dataPartJson.request_id
                                }
                                callback('data', resData);
                            }
                        }
                    });
                });
            });

            res.on('end', async () => {
                callback('end', resData);
                await prisma.t_llm_log.create({
                    data: {
                        gmt_create: new Date(),
                        input_tokens: resData.usage.input_tokens,
                        output_tokens: resData.usage.output_tokens,
                        total_tokens: resData.usage.total_tokens,
                        prompt: prompt,
                        content: resData.output.text,
                        response: JSON.stringify(resData),
                        status: 'success',
                        request_id: resData.request_id,
                        chat_request_id: context.requestId
                    }
                });
                logger.info("chat qwen successful, resp text = %s", resData.output.text);
                resolve(resData);
            });
        });

        req.on('error', async (e) => {
            logger.error('qwen chat error, errMsg: %s', e);
            callback('error', e);
            resData.err = e.message;
            await prisma.t_llm_log.create({
                data: {
                    gmt_create: new Date(),
                    input_tokens: resData.usage.input_tokens,
                    output_tokens: resData.usage.output_tokens,
                    total_tokens: resData.usage.total_tokens,
                    prompt: prompt,
                    content: resData.output.text,
                    response: JSON.stringify(resData),
                    status: 'failed',
                    error: resData.err,
                    request_id: resData.request_id,
                    chat_request_id: context.requestId
                }
            });
            reject(e);
        });

        req.write(data);
        req.end();
    });
}